# -*- coding: utf-8 -*-
"""
@Time    : 2024/7/19 15:24 
@Author  : ZhangShenao 
@File    : 2.带过滤条件的相似性检索.py
@Desc    : 带过滤条件的相似性检索

向量存储支持存储文档的元数据,如文档的作者、创建时间、版本号等等
LangChain在Faiss的基础上,封装了条件过滤功能,可以基于metadata中的属性进行过滤
"""
import dotenv
from langchain_community.embeddings import ZhipuAIEmbeddings
from langchain_community.vectorstores import FAISS

# 加载环境变量
dotenv.load_dotenv()

# 创建ZhipuAIEmbeddings
# 使用embedding-2这个Embedding模型
embeddings = ZhipuAIEmbeddings(model='embedding-2')

# 构造文本列表
texts = [
    '笨笨是一只很喜欢睡觉的猫咪',
    '我喜欢在夜晚听音乐，这让我感到放松。',
    '猫咪在窗台上打盹，看起来非常可爱。',
    '学习新技能是每个人都应该追求的目标。',
    '我最喜欢的食物是意大利面，尤其是番茄酱的那种。',
    '昨晚我做了一个奇怪的梦，梦见自己在太空飞行。',
    '我的手机突然关机了，让我有些焦虑。',
    '阅读是我每天都会做的事情，我觉得很充实。',
    '他们一起计划了一次周末的野餐，希望天气能好。',
    '我的狗喜欢追逐球，看起来非常开心。',
]

# 构造文档的元数据
metadatas = [
    {'page': 1},
    {'page': 2},
    {'page': 3},
    {'page': 4},
    {'page': 5},
    {'page': 6},
    {'page': 7},
    {'page': 8},
    {'page': 9},
    {'page': 10},
]

# 从文本导入Faiss向量数据库
faiss = FAISS.from_texts(texts=texts,  # 指定文本内容
                         metadatas=metadatas,  # 指定文档元数据
                         embedding=embeddings  # 指定Embedding模型
                         )

# 带有得分的相似度检索,返回的得分是欧氏距离,距离越短表示越相关
similarity_search_with_score_result = faiss.similarity_search_with_score(
    query='我养了一只猫，它叫笨笨',  # 指定检索关键词
    k=4,  # 指定返回的数据条数
    filter=lambda x: x['page'] > 5,  # 指定过滤条件,只返回page>5的文档
)
print(f'similarity_search_with_score_result: {similarity_search_with_score_result}')
